# main.py
from environment import MultiUAVEnv
import matplotlib.pyplot as plt

plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False


def plot_environment(env):
    # 获取网格大小
    grid_size = env.grid_size
    
    # 创建图形
    fig, ax = plt.subplots()
    
    # 绘制地面用户
    served_label_added = False
    unserved_label_added = False
    for gu in env.ground_users:
        if gu.served:
            if not served_label_added:
                ax.plot(gu.position[0], gu.position[1], 'bo', label='满足需求的地面用户')
                served_label_added = True
            else:
                ax.plot(gu.position[0], gu.position[1], 'bo')
        else:
            if not unserved_label_added:
                ax.plot(gu.position[0], gu.position[1], 'yo', label='还未满足的地面用户')
                unserved_label_added = True
            else:
                ax.plot(gu.position[0], gu.position[1], 'yo')
    
    # 绘制无人机及其覆盖半径
    uav_label_added = False
    for uav in env.uavs:
        if not uav_label_added:
            ax.plot(uav.position[0], uav.position[1], 'ro', label='无人机')
            uav_label_added = True
        else:
            ax.plot(uav.position[0], uav.position[1], 'ro')
        circle = plt.Circle(uav.position, 100, color='r', fill=False, linestyle='--')
        ax.add_artist(circle)
    
    # 绘制充电站
    station_label_added = False
    for station in env.charging_stations:
        if not station_label_added:
            ax.plot(station[0], station[1], 'go', label='充电站')
            station_label_added = True
        else:
            ax.plot(station[0], station[1], 'go')
    
    # 设置坐标轴范围
    ax.set_xlim(0, grid_size)
    ax.set_ylim(0, grid_size)
    
    # 添加图例
    ax.legend()
    
    # 显示图形
    plt.title('多无人机环境及覆盖半径')
    plt.xlabel('X坐标')
    plt.ylabel('Y坐标')
    plt.grid(True)
    plt.show()


def main():
    num_uavs = 3  
    energy_states = 10+1
    cruise_states = 1000*1000  # 假设网格大小为10x10，所以状态数为 100
    grid_size = 1000  # 网格大小
    episodes = 50000  # 设置训练的回合数
    num_gu = 20

    # 初始化多无人机环境
    env = MultiUAVEnv(num_uavs, energy_states, cruise_states, grid_size,num_gu)

    # 运行模拟
    for episode in range(episodes):
        print(f"Episode {episode + 1}/{episodes}")
        #plot_environment(env)
        env.run_episode()
        # 检查所有地面用户是否都已服务
        all_served = all(gu.served for gu in env.ground_users)
        if all_served:
            print("所有地面用户的需求都已满足！")
    
    
    plot_environment(env)
    
        
    print("Training completed!")

if __name__ == "__main__":
    main()
# todo:奖励函数需要修改
# todo:需要加入对于信道的建模


